# -*- coding=utf-8 -*-
"""
@author: xingwg
@license: (C) Copyright 2020-2025.
@contact: xingweiguo@chinasvt.com
@project: boya-reid
@file: base_dataset.py
@time: 2020/9/12 19:51
@desc:
"""
import os
import glog
import cv2
# from PIL import Image
from prettytable import PrettyTable
from torch.utils.data import Dataset


class BaseImageDataset(object):
    """
    ReID数据集基类
    """
    @staticmethod
    def get_dataset_info(dataset):
        """
        获取数据集的摄像头ID信息、图片ID和数量等信息
        :param dataset:
        :return:
        """
        person_ids = []
        cam_ids = []
        for _, person_id, cam_id in dataset:
            person_ids += [person_id]
            cam_ids += [cam_id]
        person_ids = set(person_ids)
        cam_ids = set(cam_ids)
        num_person_ids = len(person_ids)
        num_cam_ids = len(cam_ids)
        num_images = len(dataset)

        return num_person_ids, num_images, num_cam_ids

    def print_dataset_statistics(self, train_images, query_images, gallery_images):
        """
        统计打印数据信息
        :param train_images:
        :param query_images:
        :param gallery_images:
        :return:
        """
        num_train_person_ids, num_train_images, num_train_cam_ids  = self.get_dataset_info(train_images)
        num_query_person_ids, num_query_images, num_query_cam_ids  = self.get_dataset_info(query_images)
        num_gallery_person_ids, num_gallery_images, num_gallery_cam_ids = self.get_dataset_info(gallery_images)

        glog.info("Dataset statistics:")
        table = PrettyTable(["subset", "ids", "images", "cameras"])
        table.add_row(["train", num_train_person_ids, num_train_images, num_train_cam_ids])
        table.add_row(["query", num_query_person_ids, num_query_images, num_query_cam_ids])
        table.add_row(["gallery", num_gallery_person_ids, num_gallery_images, num_gallery_cam_ids])
        glog.info("\n{}".format(table))


class ImageDataset(Dataset):

    def __init__(self, dataset, transform=None):
        self.dataset = dataset
        self.transform = transform

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, item):
        image_path, person_id, cam_id = self.dataset[item]
        image_name = image_path.split("/")[-1]
        if not os.path.exists(image_path):
            raise IOError("{} does not exist".format(image_path))
        # img_rgb = Image.open(image_path)
        img_bgr = cv2.imread(image_path)
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        img_rgb = img_rgb / 255.0
        if self.transform is not None:
            img_rgb = self.transform(img_rgb)

        return img_rgb, person_id, cam_id, image_name
